#define FLEXIO_DEV_VER_USED FLEXIO_DEV_VER(24, 9, 0)

#include <doca_dpa_dev.h>
#include <doca_dpa_dev_buf.h>
#include <libflexio-libc/stdio.h>
#include <libflexio-libc/string.h>
#include <libflexio-dev/flexio_dev_debug.h>
#include <libflexio-dev/flexio_dev.h>
#include <libflexio-dev/flexio_dev_queue_access.h>
#include <libflexio-dev/flexio_dev_endianity.h>
#include <dpaintrin.h>
#include "DeviceShare.h"


#define LOG_LEVEL_OFF 0
#define LOG_LEVEL_ERROR 1
#define LOG_LEVEL_DEBUG 2
#define LOG_LEVEL_INFO 3

#define CUR_LOG_LEVEL LOG_LEVEL_INFO

#ifndef CUR_LOG_LEVEL
#define CUR_LOG_LEVEL LOG_LEVEL_OFF
#endif

static inline const char *__get_file_name(const char *file_name) {
    if (strrchr(file_name, '/') == NULL) {
        return file_name;
    }
    return strrchr(file_name, '/') + 1;
}

#define __FILENAME__ __get_file_name(__FILE__)

#if defined(__x86_64__) || defined(__aarch64__)
#define LOG_ERROR HOST_LOG_ERROR
#define LOG_DEBUG HOST_LOG_DEBUG
#define LOG_INFO HOST_LOG_INFO
#define PrintRaw(...) printf(__VA_ARGS__)
#define LOG_F(stream_id, ...)
#elif defined(__riscv)
#define LOG_ERROR DEV_LOG_ERROR
#define LOG_DEBUG DEV_LOG_DEBUG
#define LOG_INFO DEV_LOG_INFO
#define PrintRaw(...) flexio_dev_print(__VA_ARGS__)
#define LOG_F(stream_id, ...) flexio_dev_msg(stream_id, FLEXIO_MSG_DEV_INFO, __VA_ARGS__)
#endif

#define DEV_LOG_ERROR(...)                                                        \
    flexio_dev_print("[ERROR][%s:%d][%s]", __FILENAME__, __LINE__, __FUNCTION__); \
    flexio_dev_print(__VA_ARGS__)
#define DEV_LOG_DEBUG(...)                                                        \
    flexio_dev_print("[DEBGU][%s:%d][%s]", __FILENAME__, __LINE__, __FUNCTION__); \
    flexio_dev_print(__VA_ARGS__)
#define DEV_LOG_INFO(...)                                                        \
    flexio_dev_print("[INFO][%s:%d][%s]", __FILENAME__, __LINE__, __FUNCTION__); \
    flexio_dev_print(__VA_ARGS__)
#define HOST_LOG_ERROR(...)                                             \
    printf("[ERROR][%s:%d][%s]", __FILENAME__, __LINE__, __FUNCTION__); \
    printf(__VA_ARGS__)
#define HOST_LOG_DEBUG(...)                                             \
    printf("[DEBGU][%s:%d][%s]", __FILENAME__, __LINE__, __FUNCTION__); \
    printf(__VA_ARGS__)
#define HOST_LOG_INFO(...)                                             \
    printf("[INFO][%s:%d][%s]", __FILENAME__, __LINE__, __FUNCTION__); \
    printf(__VA_ARGS__)

#if CUR_LOG_LEVEL >= LOG_LEVEL_ERROR
#define LOG_E LOG_ERROR
#else
#define LOG_ERROR(...)
#endif

#if CUR_LOG_LEVEL >= LOG_LEVEL_DEBUG
#define LOG_D LOG_DEBUG
#else
#define LOG_D(...)
#endif

#if CUR_LOG_LEVEL >= LOG_LEVEL_INFO
#define LOG_I LOG_INFO
#else
#define LOG_I(...)
#endif

#ifndef unlikely
#define likely(x) __builtin_expect(!!(x), 1)
#define unlikely(x) __builtin_expect(!!(x), 0)
#endif

#define Assert(condition)                         \
    if (!(condition))                             \
    {                                             \
        LOG_E("Assert failed: %s\n", #condition); \
        return -1;                                \
    }

/* get ptr to host_mem */
flexio_uintptr_t get_host_buffer(uint32_t window_id, uint32_t mkey, void *haddr) {
	struct flexio_dev_thread_ctx *dtctx;
	flexio_uintptr_t host_buffer;
	flexio_dev_status_t ret;
	Assert(flexio_dev_get_thread_ctx(&dtctx) >= 0);

	ret = flexio_dev_window_config(dtctx, (uint16_t)window_id, mkey);
	Assert(ret == FLEXIO_DEV_STATUS_SUCCESS);

	ret = flexio_dev_window_ptr_acquire(dtctx, (uint64_t)haddr, &host_buffer);
	Assert(ret == FLEXIO_DEV_STATUS_SUCCESS);
	return host_buffer;
}

static struct DpaThreadInfo dpa_execute_unit[256];
size_t time_collect[256];
// int check_offset[256];

/* execute sigrid_hash element-wise */
void dpa_sigrid_hash(uint64_t *input_data, uint64_t *output_data, struct DpaThreadInfo *arg) {
    int64_t size = arg->size;
    int salt = 0;
    int num_embeddings = 10;

    for (int i = 0; i < size; ++i) {
        int64_t x = ((int64_t*)input_data)[i];
        uint64_t value = (x ^ (uint64_t)(salt));
        value = (value * 2654435761U) ^ (value >> 16);
        value = value % (uint64_t)(num_embeddings);
        ((int64_t*)output_data)[i] = value;
    }
}

/* dispatch task to 256 thread, save to dpa_execute_unit[thread_id] */
__dpa_rpc__ uint64_t thread_task_init(uint64_t data) {   
	struct DpaThreadInfo *unit_data    = (struct DpaThreadInfo *)data;
	struct DpaThreadInfo Tmp = (*unit_data);

	int batch_dim  = Tmp.size;
	int thread_num = Tmp.thread_num;
	int step = batch_dim / thread_num;
	int remainder = batch_dim % thread_num;

	for (int j = 0; j < thread_num; ++j) {
		struct DpaThreadInfo T;
		T = Tmp;

		int res = j;
		if (res > remainder) res = remainder;
		int start_offset = j * step + res;
		int end_offset = start_offset + step + (j < remainder ? 1 : 0);
		int num_elements = end_offset - start_offset;

		T.input_mem_addr  = (char*) T.input_mem_addr + start_offset * sizeof(int64_t);
        T.output_mem_addr = (char*) T.output_mem_addr + start_offset * sizeof(int64_t);
		
		T.size = num_elements;
		dpa_execute_unit[j] = T;
	}

	__dpa_thread_fence(__DPA_SYSTEM, __DPA_RW, __DPA_RW);
	__dpa_thread_memory_writeback();

	return 0;
}

/* execute sigrid_hash via docasdk */
__dpa_global__ void DocaSdkProcess(doca_dpa_dev_mmap_t mmap_handle) {
	// return;
	size_t start = __dpa_thread_cycles();
	struct DpaThreadInfo *arg = (struct DpaThreadInfo *) &dpa_execute_unit[doca_dpa_dev_thread_rank()];

	uint64_t *input_data = (uint64_t *) doca_dpa_dev_mmap_get_external_ptr(mmap_handle, (uint64_t)arg->input_mem_addr);
	uint64_t *output_data = (uint64_t *) doca_dpa_dev_mmap_get_external_ptr(mmap_handle, (uint64_t)arg->output_mem_addr);

	dpa_sigrid_hash(input_data, output_data, arg);
	__dpa_thread_window_writeback();

	size_t end = __dpa_thread_cycles();
    time_collect[doca_dpa_dev_thread_rank()] = end - start;
}

/* execute sigrid_hash via Flexio */
__dpa_rpc__ uint64_t FlexioProcess(uint64_t thread_task_id) {
	// struct flexio_dev_thread_ctx *dtctx;
	// Assert(flexio_dev_get_thread_ctx(&dtctx) >= 0);
	// uint32_t thread_id = flexio_dev_get_thread_id(dtctx);

	size_t start = __dpa_thread_cycles();
	struct DpaThreadInfo *arg = (struct DpaThreadInfo *) &dpa_execute_unit[thread_task_id];
	// check_offset[thread_task_id] = (int)thread_id;

	uint64_t *input_data  = NULL;
	uint64_t *output_data = NULL;

	output_data = (uint64_t *)get_host_buffer(arg->window_id, arg->output_mem_mkey, arg->output_mem_addr);
	input_data  = (uint64_t *)get_host_buffer(arg->window_id, arg->input_mem_mkey, arg->input_mem_addr);
	
    dpa_sigrid_hash(input_data, output_data, arg);
    __dpa_thread_window_writeback();

	size_t end = __dpa_thread_cycles();
    time_collect[thread_task_id] = end - start;

	return 0;
}

__dpa_rpc__ uint64_t get_max_thread_time(void) {
    size_t max_cycles = 0;
    for (int i = 0; i < 256; ++i) {
        if (time_collect[i] > max_cycles) max_cycles = time_collect[i];
		// printf ("cycle[%d] = %lu\n", i, time_collect[i]);
		// fflush(stdout);
    }
    return max_cycles;
}
